package ragserverengine.milvus

import io.ktor.server.application.*

import ragserverengine.milvus.core.MilvusClient
import ragserverengine.milvus.core.MilvusKey

import com.google.gson.Gson
import com.google.gson.JsonObject
import io.milvus.v2.common.DataType
import io.milvus.v2.common.IndexParam
import io.milvus.v2.client.MilvusClientV2
import io.milvus.v2.service.collection.request.AddFieldReq
import io.milvus.v2.service.collection.request.CreateCollectionReq
import io.milvus.v2.service.collection.request.DescribeCollectionReq
import io.milvus.v2.service.vector.request.InsertReq

import io.milvus.v2.service.vector.request.data.FloatVec

import io.milvus.v2.service.vector.request.SearchReq
import io.milvus.v2.service.vector.response.SearchResp


class MilvusService(private val application: Application) {

    // by lazy 延迟初始化，第一次访问时才执行
    private val milvusClient: MilvusClient by lazy {
        application.attributes[MilvusKey]
    }

    // 每次访问都调用 milvusClient.getClientV2()
    private val client: MilvusClientV2
        get() = milvusClient.getClientV2()

    // miluvs插入数据方法，当collection不存在时，调用createCollectionIfNotExistsV2创建
    suspend fun insert(request: MilvusRequest): MilvusResponse {

        val collectionName = request.collectionName
        val docs = request.docs
        val metas = request.metas
        val vectors = request.vectors
        val dim = request.vectors.firstOrNull()?.size ?: 1024   // 默认值为1024， 如果 ? 左边为null

        // 检查 Collection 是否存在
        val exists = try {
            val descReq = DescribeCollectionReq.builder()
                .collectionName(collectionName)
                .build()
            client.describeCollection(descReq)
            true
        } catch (e: Exception) {
            false
        }

        // 如果不存在则创建
        if (!exists) {
            println("Collection '$collectionName' not found. Creating...")
            createCollectionIfNotExistsV2(client, collectionName, dim)
        }

        // 插入数据
        val gson = Gson()
        val data = request.vectors.mapIndexed { index, vector ->  // mapIndexed函数作用是：遍历列表同时并获取索引
            JsonObject().apply {        // JsonObject是Gson提供的JSON容器，apply的作用相当于是：1.创建一个对象，2.做{}里的操作，3.返回它自己。在这里的意思就是，创建一个json对象，然后做{}里的操作，然后返回赋值给data

                // 向量字段, gson.toJsonTree把 List<Float>转换成 JsonArray
                add("embedding", gson.toJsonTree(vector))

                // 文本和元数据, getOrNull(index) 表示安全地取出第index个元素，超出范围则返回null
                addProperty("doc", request.docs.getOrNull(index) ?: "")
                addProperty("meta", request.metas.getOrNull(index) ?: "")
            }
        }
        // 插入请求
        val insertReq = InsertReq.builder()
            .collectionName(collectionName)
            .data(data)
            .build()

        // return, runCatching{}.getOrElse{}异常捕获的高阶函数，简化了try/catch
        // 执行插入 + 错误捕获
        return runCatching {
            val insertResp = client.insert(insertReq)
            val insertedCount = insertResp.insertCnt

            println("Inserted $insertedCount rows into '$collectionName'.")

            MilvusResponse(
                collectionName = collectionName,
                insertedCount = insertedCount,
                success = true,
                message = "Inserted successfully."
            )
        }.getOrElse { e ->
            MilvusResponse(
                collectionName = collectionName,
                insertedCount = 0,
                success = false,
                message = "Insert failed: ${e.message}"
            )
        }
    }


    // 如果不存在则创建
    private fun createCollectionIfNotExistsV2(
        client: MilvusClientV2,
        collectionName: String,
        dim: Int,
    ) {

        //  创建 Schema
        val schema = client.createSchema().apply {
            addField(idField("id"))
            addField(vectorField("embedding", dim))
            addField(varcharField("doc"))
            addField(varcharField("meta"))
        }

        // 设置索引方式
        val indexParamForldField = IndexParam.builder()
            .fieldName("id")
            .indexType(IndexParam.IndexType.STL_SORT)
            .build()

        val indexParamForVectorField = IndexParam.builder()
            .fieldName("embedding")
            .indexType(IndexParam.IndexType.IVF_FLAT)
            .metricType(IndexParam.MetricType.L2)
            .extraParams(mapOf("nlist" to 2560))
            .build()

        val indexParams = listOf(indexParamForldField, indexParamForVectorField)

        //  组装创建请求
        val request = CreateCollectionReq.builder()
            .collectionName(collectionName)
            .collectionSchema(schema)
            .indexParams(indexParams)
            .enableDynamicField(false)
            .build()

        //  创建 Collection
        client.createCollection(request)

        println("Collection '$collectionName' created successfully!")
    }

    //
    // --- 辅助函数区 ---
    //
    private fun idField(
        name: String,
        autoID: Boolean = true
    ): AddFieldReq =
        AddFieldReq.builder()
            .fieldName(name)
            .dataType(DataType.Int64)
            .isPrimaryKey(true)
            .autoID(autoID)
            .build()

    private fun vectorField(
        name: String,
        dim: Int
    ): AddFieldReq =
        AddFieldReq.builder()
            .fieldName(name)
            .dataType(DataType.FloatVector)
            .dimension(dim)
            .build()

    private fun varcharField(
        name: String,
        maxLen: Int = 1024
    ): AddFieldReq =
        AddFieldReq.builder()
            .fieldName(name)
            .dataType(DataType.VarChar)
            .maxLength(maxLen)
            .build()

    suspend fun search(request: MilvusSearchRequest): MilvusSearchResponse {

        val collectionName = request.collectionName
//        val queryVectors = request.vector
        val topK = request.topK
        val metric = request.metricType
        val outputFields = request.outputFields ?: listOf("id", "doc", "meta")

        val queryVectors = request.vector.map { FloatVec(it) }

        // 构建搜索请求
        val searchReq = SearchReq.builder()
            .collectionName(collectionName)
            .outputFields(outputFields)
            .data(queryVectors)
            .topK(topK)
            .build()

        return runCatching {
            val resp: SearchResp = client.search(searchReq)

            val resultItems = resp.searchResults.map { singleQueryResults ->
                singleQueryResults.map { hit ->
                    val idValue = when (val rawId = hit.id) {
                        is Number -> rawId.toLong()
                        is String -> rawId.toLongOrNull() ?: -1L
                        else -> -1L
                    }

                    MilvusSearchResultItem(
                        id = idValue,
                        score = hit.score,
                        doc = hit.entity?.get("doc") as? String,
                        meta = hit.entity?.get("meta") as? String
                    )
                }
            }

            MilvusSearchResponse(
                collectionName = collectionName,
                results = resultItems,
                success = true,
                message = "Search completed successfully."
            )
        }.getOrElse { e ->
            MilvusSearchResponse(
                collectionName = collectionName,
                results = emptyList(),
                success = false,
                message = "Search failed: ${e.message}"
            )
        }
    }
}
